

# api_server.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch
from fastapi.middleware.cors import CORSMiddleware

app = FastAPI()

# 允许跨域请求（Dify 需要）
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# 定义请求体模型
class RerankRequest(BaseModel):
    query: str
    documents: list[str]

# 初始化模型和分词器
model_name = "/home/llm/bge-reranker-v2-m3"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
model.eval()

# 如果使用 GPU
if torch.cuda.is_available():
    model = model.cuda()

@app.post("/rerank")
async def rerank(request: RerankRequest):
    try:
        # 生成模型输入对
        pairs = [[request.query, doc] for doc in request.documents]

        # 批量编码
        with torch.no_grad():
            inputs = tokenizer(
                pairs,
                padding=True,
                truncation=True,
                max_length=5120,
                return_tensors="pt"
            )

            # GPU 加速
            if torch.cuda.is_available():
                inputs = {k: v.cuda() for k, v in inputs.items()}

            # 模型推理
            scores = model(**inputs).logits[:, 0].float()

        # 转换为 Python 列表并归一化
        scores = torch.sigmoid(scores).cpu().tolist()

        return {"scores": scores}

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=9989)
